#!/usr/bin/python
'''
    Filename - ftps.py
    Author: Gerrit Kitts <gkitts@gmail.com>        
    Usage: ftps.py (listenPort)
    
    This package provides a very basic server that will listen on a specified port
    for incoming connections and accept binary file transfers from a remote host.   Communication is through simple streaming (TCP)
    socket.  The protocol used passes the file size as the first four bytes of the stream, 
    the basename (see os.path.basename) of the file as the next twenty bytes, and then the bytes
    representing the file to be transferred.

'''

import array
import hashlib
import socket
import struct
import os
import os.path

from sys import argv

'''
    Exception sub-class to handle invalid argument specifications
'''
class ArgumentException(BaseException):
    def __init__(self):
        self.msg = "Usage: ftps.py (remoteIP) (remotePort) (file-to-transfer)"

'''
    Attempts to read a file from the connection and write it to disk in the 'dest'
    directory.  The method attempts to read a number of bytes equal to the size value
    that was provided as the first four bytes of the connection, which can lead to 
    mismatches if the size header is corrupt or invalid.
'''
def writeFile(size, name, conn):
    print "Writing file %s with size %i" % (name, size)
    f = open(name, 'wb')
    
    bytesRemaining = size
    bytesToRead = 1000
    buf = array.array("B", [0] * 1000)
    
    # Create a hash for the MD5 checksum
    m = hashlib.md5()
    
    # The number of bytes to read was specified by the protocol, so read until
    #   we have consumed every byte that they passed.  If they specified too many
    #   bytes, we will block waiting; if they specified too few we will truncate 
    #   the file.
    while bytesRemaining > 0:
        numBytes = conn.recv_into(buf, bytesToRead)
        bytesToWrite = buf[0:numBytes]
        m.update(bytesToWrite)
        f.write(bytesToWrite)
        
        bytesRemaining = bytesRemaining - numBytes
        # If we have fewer than 1000 bytes to read, only read the amount we have remaining
        if bytesRemaining < bytesToRead:
            bytesToRead = bytesRemaining
        
    # Close the connection to the client to prevent the OS from thinking we still have the
    #   socket open.
    conn.close()
    
    # Write the MD5checksum
    print "MD5 checksum is: %s" % (m.hexdigest())

'''
    Creates a socket on all IP addresses on the port specified as a parameter, and listens
    for incoming requests.  
'''
def listenForConnections(params):
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    address = params['address']
    print "Binding socket %s:%s for address %s" % (address[0], address[1], address)
    s.bind(address)
    s.listen(1)
    
    conn, rhost = s.accept()    
    
    # Read the incoming header data.  The first block should be a 4-byte integer that 
    #  specifies the file size, followed by a 20 byte block representing the basename
    #  of the file.  
    unpacker = struct.Struct("! I 20s")
    metadata = conn.recv(unpacker.size)
    unpacked_data = unpacker.unpack(metadata)
    fsize = unpacked_data[0]
    fname = unpacked_data[1]
    
    # Write the file into the destination path
    fname = os.path.join(params['dest'], fname)    
    writeFile(fsize, fname.rstrip('\0'), conn)
    
    # Close the socket afte we recieve a single file - don't try to multiplex connections,
    #   handle more than one client or even stay alive.
    s.close()
    
'''
    Base wrapper for class functionality
'''
if __name__ == '__main__':
    params = {
        'address' : ('', 7777),
        'dest' : 'output'
    }
    if (len(argv) is not 2):
        # Exception case
        raise ArgumentException 
    else:
        params['address'] = ('', int(argv[1]))
        
    if len(argv) is 3:
        dirname = argv[2]
        params['dest'] = os.path.abspath(dirname)
        print "Writing files to %s" % (params['dest'])
    
    # If the dest path doesn't exist, create it (if possible)
    if not os.path.exists(params['dest']):
        os.mkdir(params['dest'])

    listenForConnections(params)
    
